import pickle
import torch
import torch.distributed as dist
from multiprocessing.synchronize import Event
from multiprocessing.shared_memory import SharedMemory

from nanovllm.config import Config
from nanovllm.engine.sequence import Sequence
from nanovllm.models.qwen3 import Qwen3ForCausalLM
from nanovllm.layers.sampler import Sampler
from nanovllm.utils.context import set_context, get_context, reset_context
from nanovllm.utils.loader import load_model


class ModelRunner:
    """
    0-分布式环境初始化与管理
    1-模型加载初始化，推理
    2-分布式协调
    3-采样器管理，KV 缓存分配与管理
    3-CUDA 图捕获与推理加速
    4-序列预处理与推理

    """
    def __init__(self, config: Config, rank: int, event: Event | list[Event]):
        """
        初始化分布式进程组，使用NCCL后端和TCP通信
        0号进程是主进程，其他进程是工作进程
        主进程通过共享内存（SharedMemory）发送任务，并通过事件通知工作进程开始处理任务
        """
        self.config = config
        hf_config = config.hf_config
        self.block_size = config.kvcache_block_size
        self.enforce_eager = config.enforce_eager
        self.world_size = config.tensor_parallel_size
        self.rank = rank
        self.event = event
        #初始化分布式进程，使用NCCL后端和TCP通信
        dist.init_process_group("nccl", "tcp://localhost:2333", world_size=self.world_size, rank=rank)
        # 当前进程绑定到指定的 GPU 设备上
        torch.cuda.set_device(rank)
        # 记录默认数据类型
        # 必须在第2行之前执行，否则就无法获取到真正的原始默认数据类型
        default_dtype = torch.get_default_dtype()
        torch.set_default_dtype(hf_config.torch_dtype) # 设置模型数据类型
        torch.set_default_device("cuda")# 设置默认设备为cuda
        self.model = Qwen3ForCausalLM(hf_config)# 构建模型
        # 加载权重
        load_model(self.model, config.model)
        # 构建采样器 负责从模型输出的概率分布中高效地采样下一个要生成的token
        """
        它决定了模型如何从预测的概率分布中选择下一个要生成的token。
        1-决定生成内容
        2-控制生成多样性
        3-平衡质量与创造性
        为什么需要采样？
        不能简单地选择概率最高的token（即贪心解码），而需要采样呢？这有几个重要原因：
        1- 避免生成单调重复内容：如果始终选择概率最高的token，生成的文本往往会变得重复、缺乏变化
        2- 增加生成的多样性：采样允许模型在生成过程中引入一定的随机性，从而产生更丰富、更多样的输出
        3- 模拟人类语言的不确定性：人类语言本身就包含不确定性和多样性，采样机制可以更好地模拟这种特性
        4- 探索更多可能性：通过采样，模型可以探索一些概率不是最高但仍然合理的选择，可能会带来更自然、更有创意的结果
        """
        self.sampler = Sampler()
        # 预热模型，用尽量大的batch和token计算一次
        self.warmup_model()
        # 分配KV缓存，计算可用的KV缓存块数量，并且绑定到每个层的k和v
        self.allocate_kv_cache()
        if not self.enforce_eager:
            self.capture_cudagraph()# 捕获CUDA图以加速推理，原理是预先加载了计算图，但是要预存计算图会占用部分的显存
        torch.set_default_device("cpu")# 恢复默认设备为cpu
        torch.set_default_dtype(default_dtype)# 恢复默认数据类型
        # 多进程共享内存和同步
        """
        共享内存区域创建与连接：
        主进程（rank=0）负责创建一个命名的共享内存区域
        工作进程（rank>0）通过相同的名称连接到这个共享内存区域
        这种方式避免了进程间通过序列化/反序列化拷贝大量数据的开销
        """
        if self.world_size > 1:
            if rank == 0:
                # 主进程创建共享内存
                self.shm = SharedMemory(name="nanovllm", create=True, size=2**20)
                dist.barrier()# 同步
            else:
                dist.barrier()  # 等待主进程
                self.shm = SharedMemory(name="nanovllm") # 其他进程连接共享内存
                self.loop()# 子进程进入循环等待任务

    def exit(self):
        if self.world_size > 1:
            self.shm.close()
            dist.barrier()
            if self.rank == 0:
                self.shm.unlink()
        if not self.enforce_eager:
            del self.graphs, self.graph_pool
        torch.cuda.synchronize()
        dist.destroy_process_group()

    def loop(self):
        while True:
            method_name, args = self.read_shm()
            self.call(method_name, *args)
            if method_name == "exit":
                break

    def read_shm(self):
        # 环境验证：assert self.world_size > 1 and self.rank > 0 确保此方法只在多进程环境中的工作进程执行
        assert self.world_size > 1 and self.rank > 0
        # 让工作进程阻塞等待，直到主进程设置事件信号
        self.event.wait()
        # 从共享内存的前4字节读取数据长度，这里使用小端序（"little"）解析4字节整数
        n = int.from_bytes(self.shm.buf[0:4], "little")
        # 使用pickle.loads() 将字节数据转换回Python对象
        method_name, *args = pickle.loads(self.shm.buf[4:n+4])
        # 重置事件信号，为下一次通信做准备
        self.event.clear()
        return method_name, args
    # 并写入共享内存
    def write_shm(self, method_name, *args):
        # 环境验证：assert self.world_size > 1 and self.rank > 0 确保此方法只在多进程环境中的工作进程执行
        assert self.world_size > 1 and self.rank == 0
        data = pickle.dumps([method_name, *args])
        n = len(data)
        self.shm.buf[0:4] = n.to_bytes(4, "little")
        self.shm.buf[4:n+4] = data
        # 通知工作进程数据已准备好
        for event in self.event:
            event.set()

    def call(self, method_name, *args):
        if self.world_size > 1 and self.rank == 0:
            self.write_shm(method_name, *args)
        method = getattr(self, method_name, None)
        return method(*args)

    """
    1. 初始化GPU资源 -> 将模型权重从CPU传输到GPU,CUDA内核加载：触发PyTorch加载相关的CUDA内核到GPU,缓存分配：为模型各层的中间计算结果分配缓冲区
    2. 性能优化 
      PyTorch会在首次运行后优化计算图，提高后续推理速度
      预热可以帮助系统了解模型在最大负载下的内存需求，为后续KV缓存分配提供依据
    3.  内存管理
    准确测量内存峰值，合理分配KV缓存，避免内存碎片
    4.稳定性保障 
    提前发现问题：预热过程可以提前发现内存不足、配置错误等问题
    系统状态稳定：通过预热让系统达到稳定状态，减少实际服务时的性能波动
    """
    def warmup_model(self):
        # 清空CUDA设备上的缓存内存，在预热前确保GPU内存尽可能干净，避免之前操作残留的内存占用影响预热过程和内存统计
        torch.cuda.empty_cache()
        # 重置CUDA内存峰值统计计数器，确保后续测量的内存峰值是预热过程中真实使用的内存量，为后续的KV缓存分配提供准确依据
        torch.cuda.reset_peak_memory_stats()
        # 系统支持的最大批处理token数量 ，模型支持的最大序列长度 为计算预热使用的序列数量做准备
        max_num_batched_tokens, max_model_len = self.config.max_num_batched_tokens, self.config.max_model_len
        # 计算预热时使用的序列数量 在不超过总令牌限制的前提下，单次批量推理最多能容纳的 “最长序列数量”
        num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs)# //向下整除
        # 创建预热用的序列列表
        seqs = [Sequence([0] * max_model_len) for _ in range(num_seqs)]
        self.run(seqs, True)
        """
        # 释放临时计算内存：清除预热过程中产生的中间计算结果和临时张量
        # 保留模型权重和持久化结构：不会清除模型权重、KV缓存等持久化内存分配
        # 优化内存使用效率：为实际推理请求准备更充足的可用内存空间
        PyTorch 内存管理机制-只会释放未被 active tensors 引用的缓存内存，不会影响正在使用的内存数据。这是理解两次调用不会清除有用数据的关键：
        """
        # 释放临时计算内存：清除预热过程中产生的中间计算结果和临时张量
        # 保留模型权重和持久化结构：不会清除模型权重、KV缓存等持久化内存分配
        # 优化内存使用效率：为实际推理请求准备更充足的可用内存空间
        torch.cuda.empty_cache()

    def allocate_kv_cache(self):
        config = self.config
        hf_config = config.hf_config
        # 获取当前GPU内存信息
        free, total = torch.cuda.mem_get_info()
        used = total - free
        # 获取当前内存使用峰值和当前内存使用量
        peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
        current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
        # 计算每个设备上的KV头数量（分布式环境下的分片）
        num_kv_heads = hf_config.num_key_value_heads // self.world_size
        # 计算可用的KV缓存块数量
        block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * hf_config.head_dim * hf_config.torch_dtype.itemsize
        # 计算可用的KV缓存块数量
        config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes
        assert config.num_kvcache_blocks > 0
        # 分配KV缓存
        """
        这个2用于同时存储两类缓存数据：Key缓存和Value缓存。
        在Transformer架构的自注意力机制中，每个注意力头需要计算三类张量：Query(Q)、Key(K)和Value(V)。为了加速自回归生成过程，模型会缓存之前计算过的K和V张量，避免重复计算。
        """
        self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, hf_config.head_dim)
        layer_id = 0
        for module in self.model.modules():
            # 将KV缓存绑定到模型各层
            if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
                module.k_cache = self.kv_cache[0, layer_id] # 绑定Key缓存
                module.v_cache = self.kv_cache[1, layer_id] # 绑定Value缓存
                layer_id += 1

    def prepare_block_tables(self, seqs: list[Sequence]):
        max_len = max(len(seq.block_table) for seq in seqs)
        block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs]
        block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
        return block_tables

    # 处理完整的输入序列，为首次推理做准备
    def prepare_prefill(self, seqs: list[Sequence]):
        input_ids = []
        positions = []
        cu_seqlens_q = [0]
        cu_seqlens_k = [0]
        max_seqlen_q = 0
        max_seqlen_k = 0
        slot_mapping = []
        block_tables = None
        """
        # 准备输入ID、位置编码、序列长度等信息
        # 处理前缀缓存(prefix cache)优化
        # 设置推理上下文
        """
        for seq in seqs:
            seqlen = len(seq)
            input_ids.extend(seq[seq.num_cached_tokens:])
            positions.extend(list(range(seq.num_cached_tokens, seqlen)))
            seqlen_q = seqlen - seq.num_cached_tokens
            seqlen_k = seqlen
            cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q)
            cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k)
            max_seqlen_q = max(seqlen_q, max_seqlen_q)
            max_seqlen_k = max(seqlen_k, max_seqlen_k)
            if not seq.block_table:    # warmup
                continue
            for i in range(seq.num_cached_blocks, seq.num_blocks):
                start = seq.block_table[i] * self.block_size
                if i != seq.num_blocks - 1:
                    end = start + self.block_size
                else:
                    end = start + seq.last_block_num_tokens 
                slot_mapping.extend(list(range(start, end)))
        if cu_seqlens_k[-1] > cu_seqlens_q[-1]:    # prefix cache
            block_tables = self.prepare_block_tables(seqs)
        input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
        positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
        cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
        cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
        slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
        set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables)
        return input_ids, positions

    # 只处理最后一个token，用于自回归生成过程
    def prepare_decode(self, seqs: list[Sequence]):
        input_ids = []
        positions = []
        slot_mapping = []
        context_lens = []
        """
        # 仅提取每个序列的最后一个token和位置信息
        # 准备KV缓存槽位映射和上下文长度
        # 设置推理上下文
        """
        for seq in seqs:
            input_ids.append(seq.last_token)
            positions.append(len(seq) - 1)
            context_lens.append(len(seq))
            slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens  - 1)
        input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
        positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
        slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
        context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
        block_tables = self.prepare_block_tables(seqs)
        set_context(False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables)
        return input_ids, positions

    def prepare_sample(self, seqs: list[Sequence]):
        temperatures = []
        for seq in seqs:
            temperatures.append(seq.temperature)
        temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True)
        return temperatures

    @torch.inference_mode()
    def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool):
        if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
            return self.model.compute_logits(self.model(input_ids, positions))
        else:#使用CUDA图加速推理
            bs = input_ids.size(0)
            context = get_context()
            graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]
            graph_vars = self.graph_vars
            graph_vars["input_ids"][:bs] = input_ids
            graph_vars["positions"][:bs] = positions
            graph_vars["slot_mapping"].fill_(-1)
            graph_vars["slot_mapping"][:bs] = context.slot_mapping
            graph_vars["context_lens"].zero_()
            graph_vars["context_lens"][:bs] = context.context_lens
            graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
            graph.replay()
            return self.model.compute_logits(graph_vars["outputs"][:bs])

    """
    负责处理输入序列并生成下一个token
    预填充阶段(prefill)和解码阶段(decode)，并在张量并行环境下协调不同进程的工作。
    """
    def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:
        # input_ids(输入token的ID)和positions(每个 token 的位置编码信息)
        input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs)
        # 仅在主进程(rank=0)中准备采样参数，因为只有主进程负责最终的token采样
        temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
        # 执行实际的模型推理计算，返回模型输出的logits(未归一化的概率分布)
        # 代表模型对词汇表中每个可能token的预测分数。
        logits = self.run_model(input_ids, positions, is_prefill)
        # 仅在主进程中执行token采样，将模型输出的logits转换为具体的token ID
        token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None
        # 清除当前推理上下文信息
        reset_context()
        return token_ids

    @torch.inference_mode()
    def capture_cudagraph(self):
        config = self.config
        hf_config = config.hf_config
        max_bs = min(self.config.max_num_seqs, 512)
        max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size
        input_ids = torch.zeros(max_bs, dtype=torch.int64)
        positions = torch.zeros(max_bs, dtype=torch.int64)
        slot_mapping = torch.zeros(max_bs, dtype=torch.int32)
        context_lens = torch.zeros(max_bs, dtype=torch.int32)
        block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32)
        outputs = torch.zeros(max_bs, hf_config.hidden_size)
        self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
        self.graphs = {}
        self.graph_pool = None

        for bs in reversed(self.graph_bs):
            graph = torch.cuda.CUDAGraph()
            set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs])
            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup
            with torch.cuda.graph(graph, self.graph_pool):
                outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # capture
            if self.graph_pool is None:
                self.graph_pool = graph.pool()
            self.graphs[bs] = graph
            torch.cuda.synchronize()
            reset_context()

        self.graph_vars = dict(
            input_ids=input_ids,
            positions=positions,
            slot_mapping=slot_mapping,
            context_lens=context_lens,
            block_tables=block_tables,
            outputs=outputs,
        )
